package ai.moran.demo.common.es;


import ai.moran.demo.utils.JsonUtility;
import ai.moran.demo.utils.LogUtility;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import io.searchbox.client.JestClient;
import io.searchbox.client.JestClientFactory;
import io.searchbox.client.JestResult;
import io.searchbox.core.ClearScroll;
import io.searchbox.core.Delete;
import io.searchbox.core.Get;
import io.searchbox.core.Index;
import io.searchbox.core.Search;
import io.searchbox.core.SearchResult;
import io.searchbox.core.Update;
import io.searchbox.core.SearchResult.Hit;
import io.searchbox.core.search.sort.Sort;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.commons.beanutils.BeanUtils;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class EsJestService {
    private String host;
    private String port;
    private boolean isMultiThreaded;
    private int maxTotalConnectionPerRoute;
    private int maxTotalConnection;
    private int readTimeout;
    private JestClient jestClient;
    private Logger logger;

    private EsJestService(String host, String port, boolean isMultiThreaded, int maxTotalConnectionPerRoute, int maxTotalConnection, int readTimeout, Logger logger) {
        this.host = host;
        this.port = port;
        this.isMultiThreaded = isMultiThreaded;
        this.maxTotalConnectionPerRoute = maxTotalConnectionPerRoute;
        this.maxTotalConnection = maxTotalConnection;
        this.readTimeout = readTimeout;
        JestClientFactory factory = new JestClientFactory();
        factory.setHttpClientConfig(((io.searchbox.client.config.HttpClientConfig.Builder)((io.searchbox.client.config.HttpClientConfig.Builder)(new io.searchbox.client.config.HttpClientConfig.Builder("http://" + this.host.trim() + ":" + this.port.trim())).multiThreaded(this.isMultiThreaded)).defaultMaxTotalConnectionPerRoute(this.maxTotalConnectionPerRoute).maxTotalConnection(this.maxTotalConnection).readTimeout(this.readTimeout)).build());
        this.jestClient = factory.getObject();
        if (logger != null) {
            this.logger = logger;
        } else {
            this.logger = LoggerFactory.getLogger(EsJestService.class);
        }

    }

    public boolean jestResultSucceeded(JestResult jestResult, String index, String type, String id) {
        if (jestResult == null) {
            logger.error("JestResult is null", new Object[0]);
            return false;
        } else {
            if (!jestResult.isSucceeded()) {
                this.logger.error("ES Operation Failed: {} for index: {}, type: {}, id: {}", new Object[]{jestResult.getErrorMessage(), index == null ? "" : index, type == null ? "" : type, id == null ? "" : id});
            }

            return jestResult.isSucceeded();
        }
    }

    public JestClient getJestClient() {
        return this.jestClient;
    }

    public Logger getLogger() {
        return this.logger;
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, List<String> includeSources, Map<String, String> andConditions) throws IOException {
        return this.searchData(index, type, sourceType, includeSources, andConditions, 1, (String)null, (SortOrder)null, (QueryBuilder[])(new ArrayList()).toArray(new QueryBuilder[0]));
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, Map<String, String> andConditions) throws IOException {
        return this.searchData(index, type, sourceType, andConditions, 100, (String)null, (SortOrder)null);
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, Map<String, String> andConditions, int size) throws IOException {
        return this.searchData(index, type, sourceType, andConditions, size, (String)null, (SortOrder)null);
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, Map<String, String> andConditions, int size, String sortKey, SortOrder sortOrder, QueryBuilder... queryBuilders) throws IOException {
        return this.searchData(index, type, sourceType, (List)null, andConditions, size, sortKey, sortOrder, queryBuilders);
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, List<String> includeSources, Map<String, String> andConditions, int size, String sortKey, SortOrder sortOrder, QueryBuilder... queryBuilders) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        if (andConditions != null && andConditions.size() > 0) {
            andConditions.entrySet().stream().forEach((e) -> {
                boolQueryBuilder.must(QueryBuilders.termQuery(((String)e.getKey()).trim(), ((String)e.getValue()).trim()));
            });
        }

        if (queryBuilders != null) {
            QueryBuilder[] var12 = queryBuilders;
            int var13 = queryBuilders.length;

            for(int var14 = 0; var14 < var13; ++var14) {
                QueryBuilder queryBuilder = var12[var14];
                boolQueryBuilder.must(queryBuilder);
            }
        }

        if (includeSources != null && includeSources.size() > 0) {
            searchSourceBuilder.fetchSource((String[])includeSources.toArray(new String[0]), (String[])null);
        }

        searchSourceBuilder.query(boolQueryBuilder);
        searchSourceBuilder.size(size);
        if (StringUtils.isNotEmpty(sortKey)) {
            if (sortOrder != null) {
                searchSourceBuilder.sort(sortKey, sortOrder);
            } else {
                searchSourceBuilder.sort(sortKey, SortOrder.DESC);
            }
        }

        return this.searchData(index, type, sourceType, searchSourceBuilder.toString());
    }

    public <T> List<T> searchData(String index, String type, Class<T> sourceType, String query) throws IOException {
        io.searchbox.core.Search.Builder builder = new io.searchbox.core.Search.Builder(query);
        builder.addIndex(index);
        if (StringUtils.isNotEmpty(type)) {
            builder.addType(type);
        }

        Search search = builder.build();
        SearchResult result = (SearchResult)this.jestClient.execute(search);
        if (this.jestResultSucceeded(result, index, type, (String)null) && result.getTotal() > 0L) {
            List<T> sourceDataList = new ArrayList();
            List<Hit<T, Void>> searchHists = result.getHits(sourceType, false);
            searchHists.stream().forEach((hit) -> {
                sourceDataList.add(hit.source);
            });
            return sourceDataList;
        } else {
            return new ArrayList();
        }
    }

    public <T> ScrollSearchPageResult<T> pageScrollSearch(String indexName, String typeName, Integer pageSize, Integer current, Class<T> sourceClass, String idPropertyName, Sort sort, QueryBuilder... queryBuilders) throws Exception {
        return this.pageScrollSearch(indexName, typeName, pageSize, current, sourceClass, idPropertyName, (Collection)Arrays.asList(sort), queryBuilders);
    }

    public <T> ScrollSearchPageResult<T> pageScrollSearch(String indexName, String typeName, Integer pageSize, Integer current, Class<T> sourceClass, String idPropertyName, Collection<Sort> sorts, QueryBuilder... queryBuilders) throws Exception {
        if (pageSize <= 0) {
            throw new IllegalArgumentException("pageSize must be greater than 0");
        } else if (current <= 0) {
            throw new IllegalArgumentException("current must be greater than 0");
        } else {
            ScrollSearchResult<T> scrollSearchResult = this.scrollSearch(indexName, typeName, pageSize, (String)null, sourceClass, idPropertyName, (Collection)sorts, queryBuilders);
            ScrollSearchPageResult<T> scrollSearchPageResult = new ScrollSearchPageResult();
            scrollSearchPageResult.setTotal(scrollSearchResult.getTotal());
            scrollSearchPageResult.setPageSize(pageSize);
            Integer count = 0;
            Iterator<List<T>> iterator = scrollSearchResult.iterator();
            List matchOrderLastPageData = null;

            while(iterator.hasNext()) {
                count = count + 1;
                matchOrderLastPageData = (List)iterator.next();
                if (count == current) {
                    break;
                }
            }

            scrollSearchPageResult.setCurrent(count);
            scrollSearchPageResult.setData(matchOrderLastPageData);
            return scrollSearchPageResult;
        }
    }

    public <T> ScrollSearchResult<T> scrollSearch(String indexName, String typeName, Integer scrollSize, String scrollTime, Class<T> sourceClass, QueryBuilder... queryBuilders) throws Exception {
        return this.scrollSearch(indexName, typeName, scrollSize, scrollTime, sourceClass, (String)null, queryBuilders);
    }

    public <T> ScrollSearchResult<T> scrollSearch(String indexName, String typeName, Integer scrollSize, String scrollTime, Class<T> sourceClass, String idPropertyName, Collection<Sort> sorts, QueryBuilder... queryBuilders) throws Exception {
        if (scrollSize == null) {
            scrollSize = 1000;
        }

        if (sorts == null) {
            sorts = new ArrayList(1);
        }

        if ((sorts).size() == 0) {
            (sorts).add(new Sort("_uid"));
        }

        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        if (queryBuilders.length > 0) {
            BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
            QueryBuilder[] var11 = queryBuilders;
            int var12 = queryBuilders.length;

            for(int var13 = 0; var13 < var12; ++var13) {
                QueryBuilder queryBuilder = var11[var13];
                boolQueryBuilder.must(queryBuilder);
            }

            searchSourceBuilder.query(boolQueryBuilder).size(scrollSize);
        } else {
            searchSourceBuilder.query(QueryBuilders.matchAllQuery()).size(scrollSize);
        }

        if (StringUtils.isEmpty(scrollTime)) {
            scrollTime = "1m";
        }

        Search search = ((io.searchbox.core.Search.Builder)((io.searchbox.core.Search.Builder)((io.searchbox.core.Search.Builder)(new io.searchbox.core.Search.Builder(searchSourceBuilder.toString())).addIndex(indexName)).addType(typeName)).addSort((Collection)sorts).setParameter("scroll", scrollTime)).build();
        ScrollSearchResult.Builder<T> builder = new ScrollSearchResult.Builder(this, scrollTime, indexName, typeName, sourceClass);
        builder.idPropertyName(idPropertyName);
        SearchResult result = (SearchResult)this.jestClient.execute(search);
        if (!this.jestResultSucceeded(result, indexName, typeName, (String)null)) {
            return builder.build();
        } else {
            List<T> sourceDataList = new ArrayList();
            List<Hit<T, Void>> searchHits = result.getHits(sourceClass, false);

            Hit hit;
            for(Iterator var15 = searchHits.iterator(); var15.hasNext(); sourceDataList.add((T) hit.source)) {
                hit = (Hit)var15.next();
                if (StringUtils.isNotEmpty(idPropertyName)) {
                    BeanUtils.setProperty(hit.source, idPropertyName, hit.id);
                }
            }

            Object scrollIdObj = result.getValue("_scroll_id");
            if (scrollIdObj != null && scrollIdObj instanceof String) {
                builder.isSuccess(true);
                builder.nextScrollId((String)scrollIdObj);
                builder.dataItems(sourceDataList);
                builder.idPropertyName(idPropertyName);
                builder.total(result.getTotal());
                return builder.build();
            } else {
                throw new IllegalArgumentException("illegal scroll id");
            }
        }
    }

    public <T> ScrollSearchResult<T> scrollSearch(String indexName, String typeName, Integer scrollSize, String scrollTime, Class<T> sourceClass, String idPropertyName, Sort sort, QueryBuilder... queryBuilders) throws Exception {
        return this.scrollSearch(indexName, typeName, scrollSize, scrollTime, sourceClass, idPropertyName, (Collection)Arrays.asList(sort), queryBuilders);
    }

    public <T> ScrollSearchResult<T> scrollSearch(String indexName, String typeName, Integer scrollSize, String scrollTime, Class<T> sourceClass, String idPropertyName, QueryBuilder... queryBuilders) throws Exception {
        return this.scrollSearch(indexName, typeName, scrollSize, scrollTime, sourceClass, idPropertyName, new Sort("_uid"), queryBuilders);
    }

    public void clearScroll(String indexName, String typeName, String... scrollIds) {
        try {
            String[] var4 = scrollIds;
            int var5 = scrollIds.length;

            for(int var6 = 0; var6 < var5; ++var6) {
                String scrollId = var4[var6];
                ClearScroll clearScroll = (new io.searchbox.core.ClearScroll.Builder()).addScrollId(scrollId).build();
                JestResult jestResult = this.jestClient.execute(clearScroll);
                this.jestResultSucceeded(jestResult, indexName, typeName, (String)null);
            }
        } catch (Exception var10) {
            LogUtility.logException(this.logger, var10);
        }

    }

    public boolean addEsData(String indexName, String typeName, String id, Object esData) {
        try {
            Index index = ((io.searchbox.core.Index.Builder)((io.searchbox.core.Index.Builder)((io.searchbox.core.Index.Builder)(new io.searchbox.core.Index.Builder(esData)).index(indexName)).type(typeName)).id(id)).build();
            JestResult jestResult = this.jestClient.execute(index);
            return this.jestResultSucceeded(jestResult, indexName, typeName, id);
        } catch (Exception var7) {
            LogUtility.logException(this.logger, var7);
            return false;
        }
    }

    public boolean upsert(String indexName, String typeName, String id, Object esData) {
        try {
            Index index = ((io.searchbox.core.Index.Builder)((io.searchbox.core.Index.Builder)((io.searchbox.core.Index.Builder)((io.searchbox.core.Index.Builder)(new io.searchbox.core.Index.Builder(esData)).index(indexName)).type(typeName)).id(id)).refresh(true)).build();
            JestResult jestResult = this.jestClient.execute(index);
            return this.jestResultSucceeded(jestResult, indexName, typeName, id);
        } catch (Exception var7) {
            LogUtility.logException(this.logger, var7);
            return false;
        }
    }

    public <T> boolean bulkAddEsData(String indexName, String typeName, List<T> esDataList, String idPropertyName) {
        try {
            if (esDataList != null && esDataList.size() != 0) {
                io.searchbox.core.Bulk.Builder bulkBuilder = (new io.searchbox.core.Bulk.Builder()).defaultIndex(indexName).defaultType(typeName);
                Iterator var6 = esDataList.iterator();

                while(var6.hasNext()) {
                    Object esData = var6.next();

                    try {
                        bulkBuilder.addAction(((io.searchbox.core.Index.Builder)(new io.searchbox.core.Index.Builder(esData)).id(BeanUtils.getProperty(esData, idPropertyName))).build());
                    } catch (Exception var9) {
                        LogUtility.logException(this.logger, var9);
                        return false;
                    }
                }

                JestResult jestResult = this.jestClient.execute(bulkBuilder.build());
                return this.jestResultSucceeded(jestResult, indexName, typeName, (String)null);
            } else {
                return true;
            }
        } catch (Exception var10) {
            LogUtility.logException(this.logger, var10);
            return false;
        }
    }

    private RuntimeException getException(String errorTitle, String errorMsg, String indexName, String typeName, String id) {
        String exceptionMsg;
        if (StringUtils.isNotEmpty(errorMsg)) {
            exceptionMsg = String.format(errorTitle + ", %s, indexName: %s,typeName: %s,id: %s", errorMsg, indexName, typeName, id);
        } else {
            exceptionMsg = String.format(errorTitle + ", indexName: %s,typeName: %s,id: %s", indexName, typeName, id);
        }

        throw new RuntimeException(exceptionMsg);
    }

    public <T> T getEsDataById(String indexName, String typeName, String id, Class<T> sourceType) throws Exception {
        Get get = ((io.searchbox.core.Get.Builder)(new io.searchbox.core.Get.Builder(indexName, id)).type(typeName)).build();
        JestResult jestResult = this.jestClient.execute(get);
        if (!this.jestResultSucceeded(jestResult, indexName, typeName, id)) {
            throw this.getException("get es failed", jestResult.getErrorMessage(), indexName, typeName, id);
        } else {
            return jestResult.getSourceAsObject(sourceType, false);
        }
    }

    public <T> T getEsDataByIdWith404Check(String indexName, String typeName, String id, Class<T> sourceType) throws Exception {
        Get get = ((io.searchbox.core.Get.Builder)(new io.searchbox.core.Get.Builder(indexName, id)).type(typeName)).build();
        JestResult jestResult = this.jestClient.execute(get);
        if (!this.jestResultSucceeded(jestResult, indexName, typeName, id)) {
            String errorMsg = jestResult.getErrorMessage();
            if (StringUtils.isNotEmpty(errorMsg) && errorMsg.contains("404")) {
                return null;
            } else {
                throw this.getException("get es failed", errorMsg, indexName, typeName, id);
            }
        } else {
            return jestResult.getSourceAsObject(sourceType, false);
        }
    }

    public boolean updateDocument(String indexName, String typeName, String id, Object document) {
        try {
            String documentString = "{\"doc\": " + JsonUtility.toJsonString(document) + "}";
            Update updateIndex = ((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)(new io.searchbox.core.Update.Builder(documentString)).index(indexName)).type(typeName)).id(id)).build();
            JestResult updateResult = this.jestClient.execute(updateIndex);
            return this.jestResultSucceeded(updateResult, indexName, typeName, id);
        } catch (Exception var8) {
            LogUtility.logException(this.logger, var8);
            return false;
        }
    }

    public <T> boolean upsertDoc(String indexName, String typeName, String id, T sourceObject) {
        try {
            String document = "{\"script\": { \"source\": \"ctx._source.version += 1 \", \"lang\": \"painless\"},\"upsert\": " + JsonUtility.toJsonString(sourceObject) + "}";
            Update updateIndex = ((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)(new io.searchbox.core.Update.Builder(document)).index(indexName)).type(typeName)).id(id)).build();
            JestResult updateResult = this.jestClient.execute(updateIndex);
            boolean isSuccess = this.jestResultSucceeded(updateResult, indexName, typeName, id);
            if (!isSuccess) {
                throw this.getException("upsert es doc failed", (String)null, indexName, typeName, id);
            } else {
                String resultValue = updateResult.getJsonObject().getAsJsonPrimitive("result").getAsString();
                if (StringUtils.isNotEmpty(resultValue)) {
                    if (resultValue.equals("created")) {
                        return true;
                    }

                    if (resultValue.equals("updated")) {
                        return false;
                    }
                }

                return false;
            }
        } catch (Exception var10) {
            LogUtility.logException(this.logger, var10);
            return false;
        }
    }

    public boolean deleteDocument(String index, String type, String id) throws Exception {
        Delete.Builder builder = (Delete.Builder) (new Delete.Builder(id)).index(index);
        builder.type(type);
        Delete delete = builder.build();
        JestResult result = null;
        result = this.jestClient.execute(delete);
        boolean isSuccess = this.jestResultSucceeded(result, index, type, id);
        if (!isSuccess) {
            throw this.getException("delete es doc failed", (String)null, index, type, id);
        } else {
            return true;
        }
    }

    public <T> boolean docAsUpsert(String indexName, String typeName, String id, T sourceObject) throws Exception {
        try {
            String document = "{\"doc\" :" + JsonUtility.toJsonString(sourceObject) + ",\"doc_as_upsert\" : true}";
            Update updateIndex = ((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)((io.searchbox.core.Update.Builder)(new io.searchbox.core.Update.Builder(document)).index(indexName)).type(typeName)).id(id)).setParameter("retry_on_conflict", 3)).build();
            JestResult updateResult = this.jestClient.execute(updateIndex);
            boolean isSuccess = this.jestResultSucceeded(updateResult, indexName, typeName, id);
            if (!isSuccess) {
                throw this.getException("upsert doc failed", updateResult.getErrorMessage(), indexName, typeName, id);
            } else {
                String resultValue = updateResult.getJsonObject().getAsJsonPrimitive("result").getAsString();
                if (StringUtils.isNotEmpty(resultValue)) {
                    if (resultValue.equals("created")) {
                        return true;
                    }

                    if (resultValue.equals("updated")) {
                        return true;
                    }

                    if (resultValue.equals("noop")) {
                        return true;
                    }
                }

                return false;
            }
        } catch (Exception var10) {
            LogUtility.logException(this.logger, var10);
            return false;
        }
    }

    public Map strictTermMatch(Map<String, String> conditions, String indexName, String indexType) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        conditions.entrySet().stream().forEach((e) -> {
            boolQueryBuilder.must(QueryBuilders.termQuery(((String)e.getKey()).trim(), ((String)e.getValue()).trim()));
        });
        searchSourceBuilder.query(boolQueryBuilder);
        io.searchbox.core.Search.Builder builder = new io.searchbox.core.Search.Builder(searchSourceBuilder.toString());
        builder.addIndex(indexName);
        builder.addType(indexType);
        Search search = builder.build();
        SearchResult result = (SearchResult)this.jestClient.execute(search);
        return (Map)(this.jestResultSucceeded(result, indexName, indexType, (String)null) && result.getTotal() > 0L ? (Map)result.getSourceAsObject(Map.class, false) : new HashMap());
    }

    public List<Hit<Map, Void>> getHitsByStrictTermMatch(Map<String, String> conditions, String indexName, String indexType) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        conditions.entrySet().stream().forEach((e) -> {
            boolQueryBuilder.must(QueryBuilders.termQuery(((String)e.getKey()).trim(), ((String)e.getValue()).trim()));
        });
        searchSourceBuilder.query(boolQueryBuilder);
        io.searchbox.core.Search.Builder builder = new io.searchbox.core.Search.Builder(searchSourceBuilder.toString());
        builder.addIndex(indexName.trim());
        builder.addType(indexType.trim());
        Search search = builder.build();
        SearchResult result = (SearchResult)this.jestClient.execute(search);
        return (List)(this.jestResultSucceeded(result, indexName, indexType, (String)null) && result.getTotal() > 0L ? result.getHits(Map.class, false) : new ArrayList());
    }

    public static class Builder {
        private String host;
        private String port;
        private boolean isMultiThreaded = true;
        private int maxTotalConnectionPerRoute = 100;
        private int maxTotalConnection = 500;
        private int readTimeout = 60000;
        private Logger logger;

        public Builder(String host, String port) {
            this.host = host;
            this.port = port;
        }

        public EsJestService.Builder isMultiThreaded(boolean isMultiThreaded) {
            this.isMultiThreaded = isMultiThreaded;
            return this;
        }

        public EsJestService.Builder maxTotalConnectionPerRoute(int maxTotalConnectionPerRoute) {
            this.maxTotalConnectionPerRoute = maxTotalConnectionPerRoute;
            return this;
        }

        public EsJestService.Builder totalConnection(int totalConnection) {
            this.maxTotalConnection = totalConnection;
            return this;
        }

        public EsJestService.Builder readTimeout(int readTimeout) {
            this.readTimeout = readTimeout;
            return this;
        }

        public EsJestService.Builder logger(Logger logger) {
            this.logger = logger;
            return this;
        }

        public EsJestService build() {
            return new EsJestService(this.host, this.port, this.isMultiThreaded, this.maxTotalConnectionPerRoute, this.maxTotalConnection, this.readTimeout, this.logger);
        }
    }
}

